"""
任务描述：使用PyTorch实现一个卷积神经网络（CNN），对MNIST数据集进行手写数字分类。要求模型包含至少2个卷积层、1个池化层和2个全连接层，并完成训练与测试。
具体要求：
1.	模型定义（8分）
o	定义一个名为SimpleCNN的类，继承nn.Module，包含以下层：
	第1卷积层：输入通道1，输出通道16，卷积核大小3x3， stride=1， padding=1；
	ReLU激活函数 + 最大池化层（2x2，stride=2）；
	第2卷积层：输入通道16，输出通道32，卷积核大小3x3， stride=1， padding=1；
	ReLU激活函数 + 最大池化层（2x2，stride=2）；
	全连接层1：输入维度需计算（提示：MNIST图像尺寸28x28），输出维度128；
	ReLU激活函数 + Dropout层（ dropout rate=0.5）；
	全连接层2：输入维度128，输出维度10（对应10个类别）。
o	正确实现forward方法，返回未经过softmax的原始输出（logits）。
2.	数据加载（4分）
o	使用torchvision.datasets.MNIST加载训练集和测试集，要求：
	训练集：随机裁剪（28x28）、归一化（均值0.1307，标准差0.3081）、随机水平翻转；
	测试集：仅归一化；
	使用DataLoader，批次大小64，训练集开启打乱（shuffle=True）。
3.	训练过程（5分）
o	损失函数：交叉熵损失（nn.CrossEntropyLoss）；
o	优化器：Adam，学习率0.001；
o	训练轮次：5轮，每轮打印训练损失（保留4位小数）；
o	禁止使用预训练模型或迁移学习。
4.	模型评估（3分）
o	在测试集上计算分类准确率（保留2位小数）；
o	输出测试集准确率结果（如：Test Accuracy: 98.56%）。

"""
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 1. 模型定义
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # Conv Layer 1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # Conv Layer 2
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        # FC Layers: 28x28 -> 14x14 -> 7x7. Flattened size = 32 * 7 * 7
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # Conv1 + ReLU + Pool
        x = self.pool(torch.relu(self.conv2(x)))  # Conv2 + ReLU + Pool
        x = x.view(-1, 32 * 7 * 7)                # Flatten
        x = torch.relu(self.fc1(x))               # FC1 + ReLU
        x = self.dropout(x)                       # Dropout
        x = self.fc2(x)                           # FC2 (Logits)
        return x

# 2. 数据加载
transform_train = transforms.Compose([
    transforms.RandomCrop(28),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform_train)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform_test)
testloader = DataLoader(testset, batch_size=64, shuffle=False)

# 3. 训练过程
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print("开始训练...")
for epoch in range(5):
    model.train()
    running_loss = 0.0
    for images, labels in trainloader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    avg_loss = running_loss / len(trainloader)
    print(f'Epoch [{epoch+1}/5], Loss: {avg_loss:.4f}')

# 4. 模型评估
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in testloader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')